#!/usr/bin/env python3
"""
最简化的Dify API操作工具
专注：上传文档到Dify知识库
"""

import sys
import argparse
from pathlib import Path
from loguru import logger
from typing import List
from dify_client import DifyAPIClient
from config import settings

def setup_logging():
    """配置日志"""
    logger.remove()
    # 解决Windows系统编码问题
    import io
    sys.stdout = io.TextIOWrapper(sys.stdout.buffer, encoding='utf-8')
    logger.add(
        sys.stdout,
        level=settings.log_level,
        format="<green>{time:HH:mm:ss}</green> | <level>{level: <8}</level> | <level>{message}</level>",
        enqueue=True  # 异步写入，避免多线程问题
    )

def upload_files_to_knowledge_base(file_paths: List[str], knowledge_base_name: str):
    """上传文件到Dify知识库"""
    logger.info(f"准备上传文件到知识库: {knowledge_base_name}")
    
    # 验证文件是否存在或为URL
    valid_files = []
    urls = []
    
    for file_path in file_paths:
        # 检查是否为URL
        if file_path.startswith(('http://', 'https://')):
            urls.append(file_path)
            logger.info(f"识别为URL: {file_path}")
        else:
            # 检查本地文件
            path = Path(file_path)
            if path.exists() and path.is_file():
                valid_files.append(str(path.absolute()))
                logger.info(f"找到文件: {path.name}")
            else:
                logger.warning(f"文件不存在: {file_path}")
    
    if not valid_files and not urls:
        logger.error("没有找到有效的文件或URL")
        return False
    
    try:
        client = DifyAPIClient()
        logger.info(f"开始创建知识库并上传 {len(valid_files)} 个文件和 {len(urls)} 个URL...")
        
        # 处理本地文件和URL
        all_files = valid_files + urls
        result = client.upload_documents_batch(all_files, knowledge_base_name)
        
        if result['status'] == 'completed':
            logger.success(f"✅ 知识库创建成功！")
            logger.info(f"知识库ID: {result['dataset_id']}")
            logger.info(f"成功上传: {result['success']}/{result['total']} 个文件/URL")
            return True
        else:
            logger.error(f"❌ 知识库创建失败: {result.get('error', '未知错误')}")
            return False
            
    except Exception as e:
        logger.error(f"❌ 操作失败: {str(e)}")
        return False

def create_text_document(dataset_id: str, name: str, text: str):
    """通过文本创建文档"""
    try:
        client = DifyAPIClient()
        logger.info(f"创建文本文档: {name}")
        
        # 创建文档
        result = client.create_document_by_text(dataset_id, name, text)
        if not result:
            logger.error("❌ 文档创建失败")
            return False
        
        document = result['document']
        document_id = document['id']
        logger.success(f"✅ 文档创建成功，ID: {document_id}")
        
        # 轮询文档状态
        logger.info("等待文档处理完成...")
        import time
        max_attempts = 20
        attempt = 0
        
        while attempt < max_attempts:
            status_result = client.get_document_status(dataset_id, document_id)
            if status_result:
                doc_status = status_result['document']
                status = doc_status.get('indexing_status')
                display_status = doc_status.get('display_status')
                
                if status == 'completed' and display_status == 'available':
                    tokens = doc_status.get('tokens', 0)
                    segment_count = doc_status.get('segment_count', 0)
                    logger.success(f"✅ 文档处理完成！")
                    logger.info(f"   Tokens: {tokens}, 分段数: {segment_count}")
                    return True
                elif status in ['error', 'failed']:
                    logger.error(f"❌ 文档处理失败: {display_status}")
                    return False
                else:
                    logger.info(f"处理中... ({display_status})")
            
            attempt += 1
            if attempt < max_attempts:
                time.sleep(3)  # 等待3秒后重试
        
        logger.warning("⚠️ 文档处理超时，但可能仍在后台处理中")
        return False
        
    except Exception as e:
        logger.error(f"创建文本文档失败: {str(e)}")
        return False

def get_knowledge_base_detail(dataset_id: str):
    """获取知识库详情"""
    try:
        client = DifyAPIClient()
        logger.info(f"获取知识库详情: {dataset_id}")
        
        detail = client.get_dataset_detail(dataset_id)
        if detail:
            logger.success(f"✅ 知识库详情获取成功")
            logger.info(f"📋 基本信息:")
            logger.info(f"  名称: {detail.get('name', 'N/A')}")
            logger.info(f"  ID: {detail.get('id', 'N/A')}")
            logger.info(f"  描述: {detail.get('description', 'N/A')}")
            logger.info(f"  权限: {detail.get('permission', 'N/A')}")
            logger.info(f"  提供者: {detail.get('provider', 'N/A')}")
            logger.info(f"  创建时间: {detail.get('created_at', 'N/A')}")
            logger.info(f"  更新时间: {detail.get('updated_at', 'N/A')}")
            
            logger.info(f"📊 统计信息:")
            logger.info(f"  文档数量: {detail.get('document_count', 0)}")
            logger.info(f"  词汇数量: {detail.get('word_count', 0)}")
            logger.info(f"  应用数量: {detail.get('app_count', 0)}")
            
            # 检索模型信息
            retrieval_model = detail.get('retrieval_model_dict', {})
            if retrieval_model:
                logger.info(f"🔍 检索模型:")
                logger.info(f"  搜索方法: {retrieval_model.get('search_method', 'N/A')}")
                logger.info(f"  重排序: {'启用' if retrieval_model.get('reranking_enable') else '禁用'}")
                logger.info(f"  Top-K: {retrieval_model.get('top_k', 'N/A')}")
                logger.info(f"  分数阈值: {'启用' if retrieval_model.get('score_threshold_enabled') else '禁用'}")
                if retrieval_model.get('score_threshold_enabled'):
                    logger.info(f"  阈值: {retrieval_model.get('score_threshold', 'N/A')}")
            
            # 嵌入模型信息
            if detail.get('embedding_model'):
                logger.info(f"🧠 嵌入模型:")
                logger.info(f"  模型: {detail.get('embedding_model')}")
                logger.info(f"  提供商: {detail.get('embedding_model_provider', 'N/A')}")
                logger.info(f"  可用状态: {'可用' if detail.get('embedding_available', False) else '不可用'}")
            
            return True
        else:
            logger.error("❌ 获取知识库详情失败")
            return False
            
    except Exception as e:
        logger.error(f"获取知识库详情失败: {str(e)}")
        return False

def list_knowledge_bases():
    """列出所有知识库"""
    try:
        client = DifyAPIClient()
        datasets = client.get_datasets()
        if datasets:
            logger.info("知识库列表:")
            for dataset in datasets:
                logger.info(f"  - {dataset['name']} (ID: {dataset['id']})")
        else:
            logger.info("没有找到知识库")
        return True
    except Exception as e:
        logger.error(f"获取知识库列表失败: {str(e)}")
        return False

def list_documents(dataset_id: str):
    """列出知识库中的所有文档"""
    try:
        client = DifyAPIClient()
        logger.info(f"获取知识库文档列表: {dataset_id}")
        
        documents = client.get_documents(dataset_id)
        if documents:
            logger.success(f"✅ 获取到 {len(documents)} 个文档")
            logger.info("📄 文档列表:")
            for doc in documents:
                logger.info(f"  - {doc.get('name', '未知文档')} (ID: {doc.get('id', 'N/A')})")
                logger.info(f"    状态: {doc.get('indexing_status', 'N/A')} ({doc.get('display_status', 'N/A')})")
                logger.info(f"    词汇数: {doc.get('word_count', 0)}, Tokens: {doc.get('tokens', 0)}")
                logger.info(f"    创建时间: {doc.get('created_at', 'N/A')}")
                logger.info("")
        else:
            logger.info("该知识库中没有文档")
        
        return True
    except Exception as e:
        logger.error(f"获取文档列表失败: {str(e)}")
        return False

def delete_document(dataset_id: str, document_id: str):
    """删除指定文档"""
    try:
        client = DifyAPIClient()
        logger.info(f"删除文档: {document_id} (知识库: {dataset_id})")
        
        result = client.delete_document(dataset_id, document_id)
        if result:
            logger.success(f"✅ 文档删除成功")
            return True
        else:
            logger.error("❌ 文档删除失败")
            return False
    except Exception as e:
        logger.error(f"删除文档失败: {str(e)}")
        return False

def update_document_by_file(dataset_id: str, document_id: str, file_path: str, name: str = None):
    """通过文件更新文档"""
    try:
        client = DifyAPIClient()
        logger.info(f"通过文件更新文档: {document_id}")
        
        result = client.update_document_by_file(dataset_id, document_id, file_path, name)
        if result:
            document = result.get('document', {})
            logger.success(f"✅ 文档更新成功")
            logger.info(f"📄 更新后的文档信息:")
            logger.info(f"  名称: {document.get('name', 'N/A')}")
            logger.info(f"  ID: {document.get('id', 'N/A')}")
            logger.info(f"  状态: {document.get('indexing_status', 'N/A')} ({document.get('display_status', 'N/A')})")
            logger.info(f"  词汇数: {document.get('word_count', 0)}, Tokens: {document.get('tokens', 0)}")
            logger.info(f"  创建时间: {document.get('created_at', 'N/A')}")
            return True
        else:
            logger.error("❌ 文档更新失败")
            return False
    except Exception as e:
        logger.error(f"更新文档失败: {str(e)}")
        return False

def update_knowledge_base(dataset_id: str, name: str = None, indexing_technique: str = None,
                         permission: str = None, embedding_model_provider: str = None,
                         embedding_model: str = None, search_method: str = None,
                         reranking_enable: bool = None, top_k: int = None,
                         score_threshold_enabled: bool = None, score_threshold: float = None):
    """更新知识库信息"""
    try:
        client = DifyAPIClient()
        logger.info(f"更新知识库: {dataset_id}")
        
        # 构建检索模型配置
        retrieval_model = {}
        if any(param is not None for param in [search_method, reranking_enable, top_k, score_threshold_enabled, score_threshold]):
            # 获取现有配置以进行增量更新
            existing_dataset = client.get_dataset_detail(dataset_id)
            if existing_dataset and 'retrieval_model' in existing_dataset:
                retrieval_model = existing_dataset['retrieval_model']

            if search_method is not None:
                retrieval_model["search_method"] = search_method
            if reranking_enable is not None:
                retrieval_model["reranking_enable"] = reranking_enable
            if top_k is not None:
                retrieval_model["top_k"] = top_k
            if score_threshold_enabled is not None:
                retrieval_model["score_threshold_enabled"] = score_threshold_enabled
            if score_threshold is not None:
                retrieval_model["score_threshold"] = score_threshold
        
        # 更新知识库
        result = client.update_dataset(
            dataset_id=dataset_id,
            name=name,
            indexing_technique=indexing_technique,
            permission=permission,
            embedding_model_provider=embedding_model_provider,
            embedding_model=embedding_model,
            retrieval_model=retrieval_model if retrieval_model else None
        )
        
        if result:
            logger.success(f"✅ 知识库更新成功")
            logger.info(f"📋 更新后的信息:")
            logger.info(f"  名称: {result.get('name', 'N/A')}")
            logger.info(f"  索引模式: {result.get('indexing_technique', 'N/A')}")
            logger.info(f"  权限: {result.get('permission', 'N/A')}")
            logger.info(f"  嵌入模型: {result.get('embedding_model', 'N/A')}")
            logger.info(f"  提供商: {result.get('embedding_model_provider', 'N/A')}")
            
            # 显示检索模型信息
            retrieval_info = result.get('retrieval_model_dict', {})
            if retrieval_info:
                logger.info(f"🔍 检索模型:")
                logger.info(f"  搜索方法: {retrieval_info.get('search_method', 'N/A')}")
                logger.info(f"  重排序: {'启用' if retrieval_info.get('reranking_enable') else '禁用'}")
                logger.info(f"  Top-K: {retrieval_info.get('top_k', 'N/A')}")
                logger.info(f"  分数阈值: {'启用' if retrieval_info.get('score_threshold_enabled') else '禁用'}")
                if retrieval_info.get('score_threshold_enabled'):
                    logger.info(f"  阈值: {retrieval_info.get('score_threshold', 'N/A')}")
            
            return True
        else:
            logger.error("❌ 知识库更新失败")
            return False
            
    except Exception as e:
        logger.error(f"更新知识库失败: {str(e)}")
        return False

def delete_knowledge_base(dataset_id: str):
    """删除知识库"""
    try:
        client = DifyAPIClient()
        logger.info(f"删除知识库: {dataset_id}")
        
        # 删除知识库
        result = client.delete_dataset(dataset_id)
        
        if result:
            logger.success(f"✅ 知识库删除成功")
            return True
        else:
            logger.error("❌ 知识库删除失败")
            return False
            
    except Exception as e:
        logger.error(f"删除知识库失败: {str(e)}")
        return False

def main():
    """主函数"""
    parser = argparse.ArgumentParser(
        description="Dify知识库API操作工具",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
 使用示例:
   # 上传文件到知识库
   python main.py upload --files doc1.pdf doc2.txt --name "我的知识库"
   
   # 上传URL到知识库
   python main.py upload --files https://example.com/doc1.pdf https://example.com/doc2.txt --name "我的知识库"
   
   # 同时上传文件和URL到知识库
   python main.py upload --files doc1.pdf https://example.com/doc2.txt --name "我的知识库"
   
   # 通过文本创建文档
   python main.py text --dataset-id "xxx" --name "测试文档" --text "这是测试内容"
   
   # 查看知识库详情
   python main.py detail --dataset-id "知识库ID"
   
   # 更新知识库
   python main.py update --dataset-id "知识库ID" --name "新名称" --search-method "keyword_search"
   
   # 删除知识库
   python main.py delete --dataset-id "知识库ID"
   
   # 列出所有知识库
   python main.py list
   
   # 文档管理
   # 列出知识库中的所有文档
   python main.py doc-list --dataset-id "知识库ID"
   
   # 删除指定文档
   python main.py doc-delete --dataset-id "知识库ID" --document-id "文档ID"
   
   # 通过文件更新文档
   python main.py doc-update-file --dataset-id "知识库ID" --document-id "文档ID" --file "文件路径" [--name "新名称"]
        """
    )
    
    subparsers = parser.add_subparsers(dest='command', help='命令')
    
    # 上传文件命令
    upload_parser = subparsers.add_parser('upload', help='上传文件到知识库')
    upload_parser.add_argument('--files', nargs='+', required=True, help='要上传的文件路径')
    upload_parser.add_argument('--name', required=True, help='知识库名称')
    
    # 创建文本文档命令
    text_parser = subparsers.add_parser('text', help='通过文本创建文档')
    text_parser.add_argument('--dataset-id', required=True, help='目标知识库ID')
    text_parser.add_argument('--name', required=True, help='文档名称')
    text_parser.add_argument('--text', required=True, help='文档文本内容')
    
    # 详情命令
    detail_parser = subparsers.add_parser('detail', help='查看知识库详情')
    detail_parser.add_argument('--dataset-id', required=True, help='知识库ID')
    
    # 更新命令
    update_parser = subparsers.add_parser('update', help='更新知识库')
    update_parser.add_argument('--dataset-id', required=True, help='知识库ID')
    update_parser.add_argument('--name', help='知识库名称')
    update_parser.add_argument('--indexing-technique', choices=['high_quality', 'economy'], help='索引模式')
    update_parser.add_argument('--permission', choices=['only_me', 'all_team_members', 'partial_members'], help='权限')
    update_parser.add_argument('--embedding-model-provider', help='嵌入模型提供商')
    update_parser.add_argument('--embedding-model', help='嵌入模型')
    update_parser.add_argument('--search-method', choices=['keyword_search', 'semantic_search', 'full_text_search', 'hybrid_search'], help='搜索方法')
    update_parser.add_argument('--reranking-enable', action='store_true', help='是否启用重排序')
    update_parser.add_argument('--top-k', type=int, help='返回结果数量')
    update_parser.add_argument('--score-threshold-enabled', action='store_true', help='是否启用分数阈值')
    update_parser.add_argument('--score-threshold', type=float, help='分数阈值')
    
    # 删除命令
    delete_parser = subparsers.add_parser('delete', help='删除知识库')
    delete_parser.add_argument('--dataset-id', required=True, help='知识库ID')
    
    # 列表命令
    list_parser = subparsers.add_parser('list', help='列出知识库')
    
    # 文档管理命令
    # 列出文档
    doc_list_parser = subparsers.add_parser('doc-list', help='列出知识库中的所有文档')
    doc_list_parser.add_argument('--dataset-id', required=True, help='知识库ID')
    
    # 删除文档
    doc_delete_parser = subparsers.add_parser('doc-delete', help='删除指定文档')
    doc_delete_parser.add_argument('--dataset-id', required=True, help='知识库ID')
    doc_delete_parser.add_argument('--document-id', required=True, help='文档ID')
    
    # 通过文件更新文档
    doc_update_file_parser = subparsers.add_parser('doc-update-file', help='通过文件更新文档')
    doc_update_file_parser.add_argument('--dataset-id', required=True, help='知识库ID')
    doc_update_file_parser.add_argument('--document-id', required=True, help='文档ID')
    doc_update_file_parser.add_argument('--file', required=True, help='要上传的文件路径')
    doc_update_file_parser.add_argument('--name', help='文档名称（可选）')
    
    args = parser.parse_args()
    
    if not args.command:
        parser.print_help()
        return 0
    
    try:
        if args.command == 'upload':
            success = upload_files_to_knowledge_base(args.files, args.name)
        elif args.command == 'text':
            success = create_text_document(args.dataset_id, args.name, args.text)
        elif args.command == 'detail':
            success = get_knowledge_base_detail(args.dataset_id)
        elif args.command == 'update':
            success = update_knowledge_base(
                dataset_id=args.dataset_id,
                name=args.name,
                indexing_technique=args.indexing_technique,
                permission=args.permission,
                embedding_model_provider=args.embedding_model_provider,
                embedding_model=args.embedding_model,
                search_method=args.search_method,
                reranking_enable=args.reranking_enable,
                top_k=args.top_k,
                score_threshold_enabled=args.score_threshold_enabled,
                score_threshold=args.score_threshold
            )
        elif args.command == 'delete':
            success = delete_knowledge_base(args.dataset_id)
        elif args.command == 'list':
            success = list_knowledge_bases()
        elif args.command == 'doc-list':
            success = list_documents(args.dataset_id)
        elif args.command == 'doc-delete':
            success = delete_document(args.dataset_id, args.document_id)
        elif args.command == 'doc-update-file':
            success = update_document_by_file(args.dataset_id, args.document_id, args.file, args.name)
        else:
            logger.error(f"未知命令: {args.command}")
            return 1
        
        return 0 if success else 1
        
    except KeyboardInterrupt:
        logger.info("用户中断操作")
        return 0
    except Exception as e:
        logger.error(f"程序错误: {str(e)}")
        return 1

if __name__ == "__main__":
    setup_logging()
    
    logger.info("Dify知识库API工具启动")
    logger.info(f"Dify服务: {settings.dify_base_url}")
    
    exit_code = main()
    sys.exit(exit_code)